from fastapi import APIRouter, WebSocket, WebSocketDisconnect
import json
import asyncio
import os
import uuid
import requests
from typing import Dict, List
from logger import logging
from openai import OpenAI
from utils.utils import get_current_date_time_weekday
from tavily import TavilyClient
from database import db, SessionLocal
from models import Conversation
from datetime import datetime

router = APIRouter(prefix="/api")

citation_prompt = ""
with open("./prompt/citation_prompt.md", "r", encoding="utf-8") as f:
    citation_prompt = f.read().strip()

# 内存中的session存储（用作缓存）
chat_sessions: Dict[str, List[dict]] = {}


def create_session() -> str:
    """创建session"""
    session_id = str(uuid.uuid4())
    chat_sessions[session_id] = []
    return session_id


def load_session_from_db(session_id: str) -> List[dict]:
    """从数据库加载session的消息历史"""
    try:
        db_session = SessionLocal()
        conversation = db_session.query(Conversation).filter(
            Conversation.session_id == session_id
        ).first()

        if conversation and conversation.messages:
            # 将数据库中的消息格式转换为内存格式
            messages = []
            for msg in conversation.messages:
                if isinstance(msg, dict) and 'role' in msg and 'content' in msg:
                    messages.append({
                        'role': msg['role'],
                        'content': msg['content']
                    })
            return messages
        return []
    except Exception as e:
        logging.error(f"从数据库加载会话失败: {e}")
        return []
    finally:
        db_session.close()


def add_message_to_session(session_id: str, role: str, content: str, save_to_db: bool = True, source: str = "external"):
    """向session中添加消息（同时更新内存和数据库）"""
    # 如果内存中没有该session，先从数据库加载
    if session_id not in chat_sessions:
        chat_sessions[session_id] = load_session_from_db(session_id)

    # 添加新消息到内存
    chat_sessions[session_id].append({
        'role': role,
        'content': content
    })

    # 保持最近的20条对话记录，避免内存过大
    if len(chat_sessions[session_id]) > 20:
        # 保留系统消息和最近的19条对话
        system_messages = [
            msg for msg in chat_sessions[session_id] if msg['role'] == 'system']
        recent_messages = [
            msg for msg in chat_sessions[session_id] if msg['role'] != 'system'][-19:]
        chat_sessions[session_id] = system_messages + recent_messages

    # 立即保存到数据库（如果需要）
    if save_to_db and role != 'system':  # 系统消息不保存到数据库
        save_conversation_to_db(session_id, source)


def get_session_messages(session_id: str) -> List[dict]:
    """获取session的消息历史（优先从内存获取，如果内存中没有则从数据库加载）"""
    # 如果内存中有，直接返回
    if session_id in chat_sessions and chat_sessions[session_id]:
        return chat_sessions[session_id]

    # 如果内存中没有，从数据库加载
    messages = load_session_from_db(session_id)
    if messages:
        # 将加载的消息缓存到内存中
        chat_sessions[session_id] = messages
        logging.info(
            f"从数据库加载会话历史，session_id: {session_id}, 消息数: {len(messages)}")
    else:
        # 如果数据库中也没有，初始化为空列表
        chat_sessions[session_id] = []

    return chat_sessions[session_id]


def save_conversation_to_db(session_id: str, source: str = "internal"):
    """保存对话历史到数据库"""
    try:
        db_session = SessionLocal()
        messages = get_session_messages(session_id)

        # 过滤掉系统消息
        user_messages = [msg for msg in messages if msg['role'] != 'system']

        if not user_messages:
            return

        # 生成对话标题（使用第一条用户消息的前20个字符）
        first_user_message = next(
            (msg for msg in user_messages if msg['role'] == 'user'), None)
        title = first_user_message['content'][:20] + "..." if first_user_message and len(
            first_user_message['content']) > 20 else first_user_message['content'] if first_user_message else "未命名对话"

        # 转换消息格式
        formatted_messages = []
        for idx, msg in enumerate(user_messages, start=1):
            formatted_messages.append({
                'id': idx,
                'role': msg['role'],
                'content': msg['content'],
                'status': 'normal'
            })

        # 检查是否已存在该session的对话记录
        existing_conversation = db_session.query(Conversation).filter(
            Conversation.session_id == session_id
        ).first()

        if existing_conversation:
            # 更新已有对话
            existing_conversation.messages = formatted_messages
            existing_conversation.updated_at = datetime.now()
            existing_conversation.title = title
            existing_conversation.source = source
        else:
            # 创建新对话记录
            conversation = Conversation(
                session_id=session_id,
                title=title,
                messages=formatted_messages,
                source=source
            )
            db_session.add(conversation)

        db_session.commit()
        logging.info(
            f"成功保存对话历史到数据库，session_id: {session_id}, source: {source}")
    except Exception as e:
        logging.error(f"保存对话历史到数据库失败: {e}")
        db_session.rollback()
    finally:
        db_session.close()


def retrieval(query: str):
    config = db.get_config()
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {config.external_knowledge_api_key}",
    }
    data = {
        "question": query,
        "similarity_threshold": 0.2, "vector_similarity_weight": 0.3, "top_k": 1024, "use_kg": False,
        "rerank_id": "BAAI/bge-reranker-v2-m3@SILICONFLOW",
        "dataset_ids": [config.external_knowledge_id],
    }
    resp = requests.post(
        "http://localhost/api/v1/retrieval",
        headers=headers,
        json=data)
    data = resp.json().get("data")
    if type(data) is bool:
        raise ValueError("知识库请求失败，请检查密钥是否正确")
    chunks = data.get("chunks", [])
    documents = [{"document_id": chunk["document_id"],
                  "image_id": chunk["image_id"],
                  "content": chunk["content"],
                  "document_keyword": chunk["document_keyword"]} for chunk in chunks]
    documents_detail = [{"id": i, "content": chunk["content"]}
                        for i, chunk in enumerate(chunks, start=1)]
    return documents, documents_detail


async def llm_task(query: str, websocket: WebSocket, session_id: str):
    """处理LLM对话和TTS任务"""
    try:
        config = db.get_config()
        client = OpenAI(
            base_url=config.ai_model_api_url or 'http://localhost:11434/v1/',
            api_key=config.ai_model_api_key or 'ollama',
        )
        logging.info(f"开始处理LLM任务, session_id: {session_id}, query: {query}")
        # 获取知识库内容
        documents, documents_detail = retrieval(query)
        logging.info(f"检索到的知识库内容长度: {len(documents)}")
        # 获取当前配置的提示词
        prompt = config.avatar_prompt or "你是一个智能助手"
        logging.info(f"当前使用的prompt: {prompt}")
        # 获取视频列表
        videos_list = db.get_all_videos()
        result_list = [
            f"视频id:{v.id}:keywords:{','.join(v.keywords)}" for v in videos_list]
        logging.info(f"视频列表: {result_list}")
        date, time, weekday = get_current_date_time_weekday()
        # 获取session的消息历史
        messages = [
            {
                'role': 'system',
                'content': f"""
                {prompt}
                日期时间:{date},{time} 星期:{weekday}
                可用视频列表:{result_list}
                以下是诺基亚知识库：
                {documents_detail}
                以上是诺基亚知识库。
                以下是需要你遵守的规则：
                {citation_prompt}
                以上是需要你遵守的规则。
                """
            }
        ]
        messages += get_session_messages(session_id)
        # 添加用户消息到session
        add_message_to_session(session_id, 'user', query)
        messages.append({'role': 'user', 'content': query +
                        "在资料引用方面尽可能附带图片资源的ID,如果没有明确播放视频的需要，无需播放视频，无需告诉用户视频资源列表。/no_think"})
        # 调用LLM
        responses = client.chat.completions.create(
            model=config.ai_model_name or 'qwen',
            messages=messages,
            stream=True,
            tools=[
                {
                    "type": "function",
                    "function": {
                        "name": "get_video_info",
                        "description": "Get video info",
                        "parameters": {
                            "type": "object",
                            "properties": {
                                "id": {"type": "integer", "description": "视频id"},
                            },
                            "required": ["id"]
                        }
                    }
                },
                {
                    "type": "function",
                    "function": {
                        "name": "search_online",
                        "description": "get some info from online",
                        "parameters": {
                            "type": "object",
                            "properties": {
                                "query": {"type": "string", "description": "搜索关键内容"},
                            },
                            "required": ["string"]
                        }
                    }
                }
            ],
            extra_body={
                "enable_thinking": False
            }
        )
        inThink = False
        tool_name = ""
        tool_content = ""
        assistant_response = ""
        await websocket.send_text(json.dumps({"data": documents, "type": "reference"}))
        for response in responses:
            delta = response.choices[0].delta
            content = delta.content
            if inThink:
                if content == "</think>":
                    inThink = False
                continue
            if content == "<think>":
                inThink = True
                continue
            elif content:
                await websocket.send_text(json.dumps({"text": content, "type": "text"}))
                assistant_response += content
            elif delta.tool_calls:
                tool_call = delta.tool_calls[0]
                if tool_call.function.name:
                    tool_name = tool_call.function.name
                if tool_call.function.arguments:
                    tool_content += tool_call.function.arguments
            await asyncio.sleep(0.05)
        if tool_name:
            logging.info(f"LLM请求调用工具: {tool_name}，参数: {tool_content}")
        # 处理工具调用
        if tool_name == 'search_online':
            API_KEY = os.getenv('TAVILY_API_KEY')
            if not API_KEY:
                logging.error("TAVILY_API_KEY未设置，无法调用在线搜索接口")
                await websocket.send_text(json.dumps({"text": "在线搜索功能未配置API_KEY", "type": "error"}))
                return
            try:
                logging.info(f"调用在线搜索接口，查询内容: {tool_content}")
                tavily_client = TavilyClient(api_key=API_KEY)
                arguments = json.loads(tool_content) if tool_content else {}
                query = arguments.get("query")
                response = tavily_client.search(
                    query=query, include_answer="advanced")
                logging.info(f"在线搜索接口返回结果: {response}")
                answer = response['answer']
                for item in answer:
                    await websocket.send_text(json.dumps({"text": item, "type": "text"}))
                    await asyncio.sleep(0.05)
            except Exception as e:
                logging.error(f"调用在线搜索接口失败: {e}")
                await websocket.send_text(json.dumps({"text": "调用在线搜索接口失败", "type": "error"}))
        elif tool_name == 'get_video_info':
            arguments = json.loads(tool_content) if tool_content else {}
            id = arguments.get("id")
            video_info = db.get_video_by_id(id)
            if video_info:
                await websocket.send_text(json.dumps({
                    "type": "video",
                    "url": video_info.url
                }))
                tips = "已经为您找到相关的视频信息。"
                for tip in tips:
                    await websocket.send_text(json.dumps({"text": tip, "type": "text"}))
                    await asyncio.sleep(0.05)
                assistant_response = tips
            else:
                tips = "很抱歉未能找到相关视频信息。"
                for tip in tips:
                    await websocket.send_text(json.dumps({"text": tip, "type": "text"}))
                    await asyncio.sleep(0.05)
                assistant_response = tips
        # 将assistant回复添加到session
        if assistant_response:
            logging.info(f"AI回复结果:{assistant_response}")
            add_message_to_session(session_id, 'assistant', assistant_response)
            # 发送推荐问题
            await recommended_question_task(websocket, query, assistant_response)

    except Exception as e:
        logging.error(f"LLM任务处理失败: {e}")
        await websocket.send_text(json.dumps({'type': 'error', 'text': str(e)}))


async def voice_task(query: str, websocket: WebSocket, session_id: str):
    """处理LLM对话和TTS任务"""
    try:
        add_message_to_session(session_id, 'assistant',
                               query, save_to_db=False)
        for response in query:
            # 发送文本消息
            await websocket.send_text(json.dumps({"text": response, "type": "text"}))
            await asyncio.sleep(0.05)

    except Exception as e:
        logging.error(f"初始化开场白任务处理失败: {e}")
        await websocket.send_text(json.dumps({'type': 'error', 'text': str(e)}))


async def recommended_question_task(websocket: WebSocket, query: str, assistant_response: str):
    try:
        config = db.get_config()
        client = OpenAI(
            base_url=config.ai_model_api_url,
            api_key=config.ai_model_api_key,
        )
        responses = client.chat.completions.create(
            model=config.ai_model_name,
            messages=[
                {
                    'role': 'user',
                    'content': f"""
                        请你根据用户的问题和你的回答,生成几个推荐问题,简洁且有趣,吸引用户继续提问。每个问题不超过15个字,避免使用否定句。
                        只需要返回问题列表,不要包含其他内容,每个问题后面加上逗号分割符,例如:
                        问题1,问题2,问题3
                        以下是输出示例1:
                        input: 用户问题:诺基亚的愿景和使命是什么?
                        AI回答:诺基亚的愿景是通过创新技术连接人们，打造一个更美好的数字世界。我们的使命是提供卓越的通信解决方案，推动全球信息交流与合作。
                        output: 诺基亚的愿景是什么,诺基亚的使命是什么,诺基亚如何推动信息交流
                        以下是输出示例2:
                        input: 用户问题:诺基亚如何推动信息交流?
                        AI回答: 诺基亚通过提供先进的通信技术和解决方案，连接全球用户，促进信息的快速传递与共享。我们致力于打造一个无缝连接的数字世界，让每个人都能轻松获取和分享信息。
                        output: 诺基亚如何连接全球用户,诺基亚的通信技术有哪些,数字世界的未来趋势

                        以下是需要你处理的内容:
                        用户问题:{query}
                        AI回答:{assistant_response}
                        以上是需要你处理的内容。/no_think
                        """
                },
            ],
            stream=False,
            extra_body={
                "enable_thinking": False
            }
        )
        if responses and len(responses.choices) > 0:
            recommended_questions = responses.choices[0].message.content.strip(
            )
            logging.info(f"推荐问题:{recommended_questions}")
            await websocket.send_text(json.dumps({
                "type": "recommendation",
                "data": recommended_questions
            }))
    except Exception as e:
        logging.error(f"推荐问题生成失败: {e}")


@router.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    """WebSocket聊天端点"""
    await websocket.accept()
    logging.info("WebSocket连接已建立")
    session_id = None
    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_text()
            logging.info(f"收到客户端消息: {data}")
            try:
                message = json.loads(data)
                # 获取或创建session
                if 'session_id' in message and message['session_id'] not in [None, ""]:
                    session_id = message['session_id']
                    logging.info(f"使用传入的session_id: {session_id}")
                else:
                    session_id = create_session()
                    logging.info(f"创建新的session_id: {session_id}")
                    await websocket.send_text(json.dumps({
                        'type': 'session_created',
                        'session_id': session_id
                    }))
                if 'text' in message:
                    # 处理聊天消息
                    if message.get('useAI', True):
                        await llm_task(message['text'], websocket, session_id)
                    else:
                        await voice_task(message['text'], websocket, session_id)
                else:
                    await websocket.send_text(json.dumps({
                        'type': 'error',
                        'text': '消息格式错误，需要包含text字段'
                    }))
            except json.JSONDecodeError:
                await websocket.send_text(json.dumps({
                    'type': 'error',
                    'text': '消息必须是有效的JSON格式'
                }))

    except WebSocketDisconnect:
        logging.info(f"WebSocket连接已断开")
    except Exception as e:
        logging.error(f"WebSocket错误: {e}")
        try:
            await websocket.send_text(json.dumps({
                'type': 'error',
                'text': f'服务器错误: {str(e)}'
            }))
        except:
            pass
